package cc.shacocloud.luanniao.core.service.ai

import cc.shacocloud.luanniao.core.utils.Json
import com.fasterxml.jackson.core.type.TypeReference
import org.springframework.util.ResourceUtils
import java.io.InputStreamReader
import java.math.BigInteger
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util.regex.Pattern
import java.util.stream.Collectors
import java.util.stream.IntStream
import java.util.stream.Stream
import kotlin.math.pow

/**
 * GPT2 的 文本使用多少 token 实现
 *
 * @author 思追(shaco)
 */
class GPT2Tokenizer(
    private val encoder: Map<String, Any>,
    private val bpe: List<String>,
) {

    companion object {

        val DEFAULT: GPT2Tokenizer = fromPretrained(
            encoderPath = "classpath:tokenizers/gpt2/encoder.json",
            vocabPath = "classpath:tokenizers/gpt2/vocab.bpe"
        )

        /**
         * 从指定的文件创建对象
         */
        fun fromPretrained(
            encoderPath: String,
            vocabPath: String,
        ): GPT2Tokenizer {

            val encoder = ResourceUtils.getURL(encoderPath).openStream().use {
                InputStreamReader(it).use { reader ->
                    Json.mapper.readValue(reader, object : TypeReference<Map<String, Any>>() {})
                }
            }
            val bpe = ResourceUtils.getURL(vocabPath).openStream().use {
                InputStreamReader(it).use { reader -> reader.readLines() }
            }

            return GPT2Tokenizer(encoder, bpe)
        }

    }

    private val decoder: Map<Any, String> = encoder.map { it.value to it.key }.toMap()

    private val cache: MutableMap<String, String> = HashMap()
    private val byte2unicode = byteToUnicode()

    private val bpeRanks: Map<Pair<String, String>, Int> by lazy {
        bpe.mapIndexed { i, it ->
            val paris = it.split(" ")
            (paris[0] to paris[1]) to i
        }.toMap()
    }

    private val pattern: Pattern =
        "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)|\\s+".toPattern()

    /**
     * 将文本内容转为 token
     */
    fun encode(text: String): List<Int> {
        val matcher = pattern.matcher(text)
        val unicodes = mutableListOf<String>()
        val bpeTokens = mutableListOf<Int>()

        while (matcher.find()) {
            val match = matcher.group()
            val unicodeBuilder = StringBuilder()
            for (b in match.toByteArray(StandardCharsets.UTF_8)) {
                unicodeBuilder.append(byte2unicode[b.toInt()])
            }
            unicodes.add(unicodeBuilder.toString())
        }

        for (token in unicodes) {
            for (bpeToken in bpe(token).split(" ".toRegex()).dropLastWhile { it.isEmpty() }.toTypedArray()) {
                val value = encoder[bpeToken]
                bpeTokens.add(if (value is Number) value.toInt() else value.toString().toInt())
            }
        }

        return bpeTokens
    }

    /**
     * 将 token 解码为 文本
     */
    fun decode(tokens: List<Int>): String {
        val textBuilder = StringBuilder()
        val byteBufferList: MutableList<String?> = ArrayList()

        for (token in tokens) {
            textBuilder.append(decoder[BigInteger.valueOf(token.toLong())])
        }
        val text = textBuilder.toString()

        for (element in text) {
            byteBufferList.add(byte2unicode[element.code])
        }

        val byteBuffer = ByteArray(byteBufferList.size)
        for (i in byteBuffer.indices) {
            var byteString = byteBufferList[i]
            if (byteString == null) {
                byteString = " "
            }
            byteBuffer[i] = byteString[0].code.toByte()
        }

        return StandardCharsets.UTF_8
            .decode(ByteBuffer.wrap(byteBuffer))
            .array()
            .joinToString(separator = "") { it.toString() }
    }

    // ------------- 内部方法

    private fun getPairs(word: List<String>): Set<Pair<String, String>> {
        val pairs = mutableSetOf<Pair<String, String>>()
        var prevCharacter = word[0]
        for (character in word.subList(1, word.size)) {
            pairs.add(prevCharacter to character)
            prevCharacter = character
        }
        return pairs
    }

    private fun byteToUnicode(): Map<Int, String> {
        val bs = Stream.of(
            IntStream.range('!'.code, '~'.code + 1).boxed(),
            IntStream.range('¡'.code, '¬'.code + 1).boxed(),
            IntStream.range('®'.code, 'ÿ'.code + 1).boxed()
        )
            .reduce { a, b -> Stream.concat(a, b) }
            .get()
            .collect(Collectors.toList())

        val cs: MutableList<Int> = bs.toMutableList()

        var n = 0
        val max = 2.0.pow(8).toInt()
        for (b in 0 until max) {
            if (!bs.contains(b)) {
                bs.add(b)
                cs.add(max + n)
                n += 1
            }
        }
        val csString = cs.map { i -> String(Character.toChars(i)) }

        val output = mutableMapOf<Int, String>()
        for (i in bs.indices) {
            output[bs[i]] = csString[i]
        }
        return output
    }

    private fun bpe(token: String): String {
        if (cache.containsKey(token)) {
            return cache[token]!!
        }

        var word = token.chars()
            .mapToObj { i: Int -> i.toChar().toString() }
            .collect(Collectors.toList())

        var pairs = getPairs(word)

        while (true) {
            var minScore = Int.MAX_VALUE
            var biGram: Pair<String, String>? = null

            for (pair in pairs) {
                if (bpeRanks.containsKey(pair)) {
                    val score = bpeRanks[pair]!!

                    if (score < minScore) {
                        minScore = score
                        biGram = pair
                    }
                }
            }

            if (biGram == null) {
                break
            }

            val first = biGram.first
            val second = biGram.second
            val newWord: MutableList<String> = ArrayList()
            var i = 0

            while (i < word.size) {
                val j = indexWithStartPosition(word, first, i)

                if (j != -1) {
                    newWord.addAll(word.subList(i, j))
                    i = j
                } else {
                    newWord.addAll(word.subList(i, word.size))
                    break
                }

                if (word[i] == first && i < word.size - 1 && word[i + 1] == second) {
                    newWord.add(first + second)
                    i += 2
                } else {
                    newWord.add(word[i])
                    i += 1
                }
            }

            word = newWord
            if (word.size == 1) {
                break
            } else {
                pairs = getPairs(word)
            }
        }

        val output = java.lang.String.join(" ", word)
        cache[token] = output
        return output
    }

    private fun <T> indexWithStartPosition(list: List<T>?, find: T, startPosition: Int): Int {
        if (list.isNullOrEmpty()) {
            return -1
        }
        for (index in startPosition until list.size) {
            if (list[index] == find) {
                return index
            }
        }
        return -1
    }


}
